'''
@Company: TWL
@Author: xue jian
@Email: xuejian@kanzhun.com
@Date: 2019-10-04 11:11:31
@LastEditTime: 2019-10-20 20:05:44
'''
import src.line as li
import numpy as np

import matplotlib.pyplot as plt
import matplotlib
from matplotlib import rc


def cm2inch(value):
    return value/2.54

##to make the label looks like the latex format.
matplotlib.rcParams['mathtext.fontset'] = 'cm'
matplotlib.rcParams['mathtext.rm'] = 'Bitstream Vera Sans'
matplotlib.rcParams['mathtext.it'] = 'Bitstream Vera Sans:italic'
matplotlib.rcParams['mathtext.bf'] = 'Bitstream Vera Sans:bold'
#############show AParticle#################################

font={'family':'Times New Roman','size':10}

plt.rc('font',**font)

class MyPlot(object):
    def __init__(self, save_path):
        self.save_path = save_path
        self.labels = []
        self.overlapf = []
        self.rand_overlap = []
        self.gap = []
        if self.save_path[-1] != '/':
            self.save_path += '/'

    def add_line(self, path, label):
        self.labels.append(label)
        line = li.Line(path, label)
        line.read()
        if self.overlapf == []:
            self.overlapf.append(line.order)
            self.rand_overlap.append(line.order)
        self.overlapf.append(line.overlapf)
        self.rand_overlap.append(line.rand_overlap)
        self.gap.append(line.gap)

    def get_overlap_fig(self, tag):
        tmp = np.array(self.overlapf)
        # print(tmp)
        plt.figure(figsize = (cm2inch(8), cm2inch(8)))
        plt.title(tag)
        plt.xscale("log")
        for i in range(1, len(self.overlapf)):
            plt.plot(tmp[0, :], tmp[i, :] ,linewidth = 0.5, label = self.labels[i-1])
        plt.legend(loc = 4)
        plt.xlim(1,1000)
        plt.ylim(0, 1)
        plt.yticks([0, 0.5, 1])
        plt.xlabel(r"$t_a$")
        plt.ylabel(r"$\langle \psi(t_a)|\psi_f\rangle$")
        plt.tick_params(top=True, right=True, direction='in')
        plt.tight_layout()
        save_file = self.save_path + tag + '.png'
        plt.savefig(save_file)


    def get_gap_fig(self, tag):
        tmp = np.array(self.gap)
        plt.figure(figsize = (cm2inch(8), cm2inch(8)))
        plt.title(tag)
        # plt.xscale("log")
        plt.plot(tmp, linewidth = 0.5)
        # plt.ylim(0, 1)
        plt.xticks(range(len(self.labels)), self.labels)
        plt.xlabel(r"$types$")
        plt.ylabel(r"$E_g$")
        plt.tick_params(top=True, right=True, direction='in')
        plt.tight_layout()
        save_file = self.save_path + tag + '.png'
        plt.savefig(save_file)

    def get_rand_overlap_fig(self, tag, row, col):
        '''
        @note:the  row means the rowth add figure, and col means the colth from 1, 2, 5 and so on.
        '''
        tmp = np.array(self.rand_overlap[row][col])
        plt.figure(figsize = (cm2inch(8), cm2inch(8)))
        plt.title(tag)
        # plt.xscale("log")
        plt.plot(tmp, linewidth = 0.5)
        plt.ylim(0, 1)
        plt.yticks([0, 0.5, 1])
        plt.xlabel(r"$N_{random}$")
        plt.ylabel(r"$\langle \psi(t_a)|\psi_f\rangle$")
        plt.tick_params(top=True, right=True, direction='in')
        plt.tight_layout()
        save_file = self.save_path + tag + '.png'
        plt.savefig(save_file)
        